import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np


def gazeto3d(gaze):
    gaze_gt = np.zeros([3])
    gaze_gt[0] = -np.cos(gaze[1]) * np.sin(gaze[0])
    gaze_gt[1] = -np.sin(gaze[1])
    gaze_gt[2] = -np.cos(gaze[1]) * np.cos(gaze[0])
    return gaze_gt


def angular(gaze, label):
    total = np.sum(gaze * label)
    return np.arccos(min(max(total / (np.linalg.norm(gaze) * np.linalg.norm(label)), -0.9999999), 0.9999999)) * 180 / np.pi


class CRLoss(nn.Module):
    def __init__(self, tau=0.07, lambda_prior=0.07):
        super().__init__()
        self.tau = tau  # 温度
        self.lambda_prior = lambda_prior    # 角度误差

    def similarity_fn(self, gaze_labels):
        """
        计算基于KL散度的相似性矩阵
        gaze_labels: (B, 2) 包含pitch和yaw角度
        """
        B = gaze_labels.shape[0]
        gaze_3d = np.array([gazeto3d(g) for g in gaze_labels.cpu().numpy()])  # 转换为3D向量
        S = np.zeros((B, B))

        for i in range(B):
            for j in range(B):
                S[i, j] = angular(gaze_3d[i], gaze_3d[j])
        print(S)
        return torch.log(self.lambda_prior / (torch.tensor(S, dtype=torch.float32, device=gaze_labels.device) + 1e-9))

    def forward(self, features, gaze_labels):
        """
        features: 投影头输出的特征 (2B, D)
        gaze_labels: 真实gaze方向 (B, 2) -> 扩展为(2B, 2)
        """
        batch_size = features.size(0) // 2

        # 数据增强后的标签扩展
        gaze_labels = torch.cat([gaze_labels, gaze_labels], dim=0)

        # 计算相似性矩阵
        S = self.similarity_fn(gaze_labels)  # (2B,2B)
        print(S.device)

        # 计算余弦相似度
        features_norm = F.normalize(features, p=2, dim=1)   # 对第一列用L2范式归一化
        sim_matrix = torch.mm(features_norm, features_norm.T) / self.tau  # (2B,2B)

        # 排除自身对比
        mask = ~torch.eye(2 * batch_size, dtype=torch.bool, device=features.device)

        # 计算分子分母
        S_pos = F.relu(S) * mask.float()
        exp_sim = torch.exp(sim_matrix) * mask.float()

        numerator = (S_pos * exp_sim).sum(dim=1)  # (2B,) 分子
        denominator = (torch.abs(S) * exp_sim).sum(dim=1)  # (2B,) 分母

        # 对比损失
        contrast_loss = -torch.log(numerator / (denominator + 1e-9)).mean()

        return contrast_loss

# 测试代码
batch_size = 8
feature_dim = 128
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
test_features = torch.randn(2 * batch_size, feature_dim) .to(device)  # 随机生成特征 (2B, D)
# 这个直接是归一化后的角度
test_gaze_labels = torch.randn(batch_size, 2).to(device)  # 随机 gaze 真实标签 (B, 2)

criterion = CRLoss(tau=0.07)
loss = criterion(test_features, test_gaze_labels)

print("Loss:", loss.item())
